import rospy
from std_msgs.msg import Float32
from controller import Robot

robot_belt = Robot()
robot_belt_name = robot_belt.getName()
timestep2 = int(robot_belt.getBasicTimeStep())
motor_belt = robot_belt.getDevice("belt_motor")
motor_belt.setPosition(float('+inf'))

motor_belt_vel = rospy.get_param("/motor_belt_init_vel")

def set_motor_belt_vel(data):
    global motor_belt_vel
    motor_belt_vel = data.data

# ==== ros init =======
rospy.init_node(robot_belt_name)
rospy.Subscriber("~set_motor_belt_vel",Float32,callback=set_motor_belt_vel)

if __name__ == "__main__":
    while robot_belt.step(timestep2) != -1 and not rospy.is_shutdown():
        motor_belt.setVelocity(motor_belt_vel)

